
#!/usr/bin/env python
# coding = utf-8

import configparser
import os
import sys
import boto3
import json
import time
import datetime
import xlsxwriter
from openpyxl import Workbook,load_workbook


class DateEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj,datetime.datetime):
            return obj.strftime("%Y-%m-%d %H:%M:%S")
        else:
            return json.JSONEncoder.default(self,obj)



def security_group():

    #####################
    #   Set the table   #
    #####################

    workbook = xlsxwriter.Workbook("Security_group_inbound.xlsx")
    worksheet = workbook.add_worksheet("inbound")
    title = [u'SecurityGroupName', u'SecurityGroupID', u'VPC_ID', u'Protocol', u'Portrange', u'Source']
    format_title = workbook.add_format()
    format_title.set_border(1)
    format_title.set_bg_color('#cccccc')

    format_text = workbook.add_format()
    format_text.set_border(1)

    worksheet.set_column('A:K',40)   #设置A到K列宽度20像素
    worksheet.write_row('A1',title,format_title)   #将title列表有A1开始横向写入，并且格式为format_title


    client = boto3.client('ec2')
    instance = client.describe_security_groups(
            )['SecurityGroups']

    if not instance:
        return None

    #################
    #   Variable    #
    #################

    length = len(instance)
    From_dict = {}
    IpProtocol_dict = {}
    CidrIp_dict = {}
    No_IpPermissions = {}
    count = 1
    count1 = 2
    data = []


    # data数据获取

    for i in range(length):
        attached_security_group_ids = []
        GroupName = instance[i]['GroupName']
        GroupId = instance[i]['GroupName']
        VpcId = instance[i]['VpcId']

        instance_id = instance[i]['IpPermissions']
        lengthid = len(instance_id)
        if instance_id:

            for k in range(lengthid):
                lengthip = len(instance_id[k]['IpRanges'])
                lengthips = len(instance_id[k]['UserIdGroupPairs'])
                if "FromPort" in instance_id[k].keys():
                    FromPort = instance_id[k]['FromPort']
                    ToPort = instance_id[k]['ToPort']
                    IpProtocol = instance_id[k]['IpProtocol']
#                    print(instance_id[k]['UserIdGroupPairs'])
                    IpProtocol_dict.setdefault(GroupName, []).append(instance_id[k]['IpProtocol'])
                    if FromPort == ToPort and instance_id[k]['IpRanges'] != []:
                        From_dict.setdefault(GroupName, []).append(FromPort)
#                        lengthip = len(instance_id[k]['IpRanges'])
                        for n in range(lengthip):
#                            GroupIds = instance_id[k]['UserIdGroupPairs'][0]['GroupId']
                            CidrIp = instance_id[k]['IpRanges'][n]['CidrIp']
#                           Description = instance_id[k]['IpRanges'][n]['Description']
                            CidrIp_dict.setdefault(GroupName,[]).append(instance_id[k]['IpRanges'][n]['CidrIp'])
                            step = len(CidrIp)
                            data.append([GroupName, GroupId, VpcId, IpProtocol, FromPort, CidrIp])
                            count += step

                    if instance_id[k]['IpRanges'] == [] and FromPort == ToPort:
                        for n in range(lengthips):
                            GroupIds = instance_id[k]['UserIdGroupPairs'][n]['GroupId']
                            data.append([GroupName, GroupId, VpcId, IpProtocol, FromPort, GroupIds])



                    if instance_id[k]['IpRanges'] != [] and FromPort != ToPort:
                        FromPort = str(FromPort)+'-'+str(ToPort)
                        for n in range(lengthip):
                            CidrIp = instance_id[k]['IpRanges'][n]['CidrIp']
                            data.append([GroupName, GroupId, VpcId, IpProtocol, FromPort, CidrIp])

                    if instance_id[k]['IpRanges'] == [] and FromPort != ToPort:
                        FromPort = str(FromPort)+'-'+str(ToPort)
                        GroupIds = instance_id[k]['UserIdGroupPairs'][0]['GroupId']
                        data.append([GroupName, GroupId, VpcId, IpProtocol, FromPort, GroupIds])

                    if instance_id[k]['UserIdGroupPairs'] != [] and FromPort != ToPort and len(instance_id[k]['UserIdGroupPairs']) > 1:
                        FromPort = str(FromPort)+'-'+str(ToPort)
                        for n in range(lengthips):
                            GroupIds = instance_id[k]['UserIdGroupPairs'][n]['GroupId']
                            data.append([GroupName, GroupId, VpcId, IpProtocol, FromPort, GroupIds])

                    if FromPort == ToPort and instance_id[k]['IpRanges'] != []:
                        for n in range(lengthips):
                            GroupIds = instance_id[k]['UserIdGroupPairs'][n]['GroupId']
                            data.append([GroupName, GroupId, VpcId, IpProtocol, FromPort, GroupIds])

                elif len(instance_id[k]) > 0:
#                    print(instance_id[k])
                    IpProtocol = 'ALL'
                    FromPort = 'ALL'
                    if instance_id[k]['IpRanges'] != []:
                        CidrIp = instance_id[k]['IpRanges'][0]['CidrIp']
                        data.append([GroupName, GroupId, VpcId, IpProtocol, FromPort, CidrIp])
                    elif instance_id[k]['UserIdGroupPairs'] != []:
                        CidrIp = instance_id[k]['UserIdGroupPairs'][0]['GroupId']
                        data.append([GroupName, GroupId, VpcId, IpProtocol, FromPort, CidrIp])
                    else:
                        CidrIp = 'ALL'
                        data.append([GroupName, GroupId, VpcId, IpProtocol, FromPort, CidrIp])

                else:
                    data.append([GroupName, GroupId, VpcId, '-', '-', '-'])

        else:
            data.append([GroupName, GroupId, VpcId, '-', '-', '-'])


    counts = 1
    for i in data:
        print(i)
        worksheet.write_row('A'+str(counts+1), i, format_text)
        counts += 1

    workbook.close()

    return data


def openpyxl_sample():
    if os.path.exists("smp_aws_sg.xlsx"):
        wb = load_workbook('smp_aws_sg.xlsx')
    else:
        wb = Workbook("smp_aws_sg.xlsx")
        wb.create_sheet('prd',0)
        wb.create_sheet('stg',1)

    sheetname = 'prd'

    if sheetname == 'prd':
        del wb[sheetname]
        wb.create_sheet('prd',0)
    if sheetname == 'stg':
        del wb[sheetname]
        wb.create_sheet('stg',1)

    ws = wb.get_sheet_by_name(sheetname)

    #if 'stg' in wb.sheetnames:
    #   ws = wb.get_sheet_by_name(sheetname)
    #else:
    #    wb.create_sheet('stg',1)
    #    ws = wb.get_sheet_by_name(sheetname)

    #del wb[sheetname]
    #del wb[sheetname]

    ws.append( ['SecurityGroupName', 'SecurityGroupID', 'VPC_ID', 'Protocol', 'Portrange', 'source'])

    data = security_group()

    #print (data)
    for i in range(len(data)):
         ws.append(data[i])

    wb.save("smp_aws_sg.xlsx")



if __name__ == '__main__':
    security_group()
#    openpyxl_sample()
